import os, sys, argparse, yaml, math
import torch
import numpy as np
import open3d as o3d
import trimesh
sys.path.append('./')
sys.path.append('./dif')
from lib.fk.FK_layer_no_delta import ShadowLayer
from dif.dif_net import DeformedImplicitField
from lib.fk.utils import o3d_read_mesh, get_4x4_matrix, R2q
from lib.tf.utils import get_grasp_hand_lineset
from lib.tf.FC_loss import FCLoss

import logging

log = logging.getLogger('trimesh')
file_handler = logging.FileHandler('trimesh_logs.log')
log.addHandler(file_handler)
log.setLevel(logging.ERROR)

init_dir = '/media/wind/Share/Projects/graspd/grasping/data/grasp_initial'

class GraspSythesis:
    def __init__(self, cat, init_path, device):
        self.cat = cat
        self.init_dir = init_path
        self.device = device
        
        self.scale = 10.0
        self.backoff_interval = 0.01
        self.record_every = 20
        self.refine_steps = 1000
        
        self.fk = ShadowLayer(device)
        
        dif_config = 'dif/configs/generate/{0}.yml'.format(cat)
        with open(os.path.join(dif_config),'r') as stream: 
            meta_params = yaml.safe_load(stream)
        dif_model = DeformedImplicitField(**meta_params)
        dif_model.load_state_dict(torch.load(meta_params['checkpoint_path']))
        dif_model.to(device)
        self.dif = dif_model
        
        idx = torch.Tensor([513]).squeeze().long().cuda()[None,...]
        self.code = self.dif.get_latent_code(idx)    
        
        self.fc = FCLoss(device)    
                
    def get_init_grasps(self, num=10, start_idx=0):
        g_list = []
        for i in range(start_idx, start_idx + num):
            grasp_config = np.load(f'{self.init_dir}/{i}.npz')
            root_base = np.zeros(7, dtype=np.float32)
            root_base[:3] = grasp_config['t']
            root_base[3:] = grasp_config['q'] # (x,y,z,w)
            hand_q = root_base
            shadow_q = grasp_config['j']
            g_list.append(np.concatenate([hand_q, shadow_q], axis=0))
        g_list = np.array(g_list, dtype=np.float32)
        return g_list

    def preprocess_grasps(self, grasp_list, backoff=False):
        root = torch.FloatTensor(grasp_list[:, :7]).to(self.device)
        joint = torch.FloatTensor(grasp_list[:, 7:]).to(self.device)
        root.requires_grad_(False)
        joint.requires_grad_(False)
        if backoff:
            out = self.fk.get_posed_pc(root, joint)
            # out = self.fk(joint)
            pcs = out.verts
            
            for i in range(pcs.size(0)):
                pc = pcs[i]
                pc *= self.scale
                com = torch.mean(pc, dim=0, keepdim=True)
                for j in range(0, 1000):
                    pc_backoff = pc + com * j * (self.scale * self.backoff_interval)
                    sdf = self.dif.inference(pc_backoff, self.code)
                    dist = torch.min(sdf)
                    if dist > 0:
                        print(f'------ backoff {j * 0.01} meters for {i+1}th grasp. ------')
                        root[i][:3] += j * self.backoff_interval * com[0]
                        # if True:
                        #     import trimesh
                        #     pc_1 = trimesh.PointCloud(pc.detach().cpu().numpy())
                        #     pc = trimesh.PointCloud(pc_backoff.detach().cpu().numpy(),colors=np.array([255,0,0,255]))
                        #     mesh = trimesh.load('dif/recon/mug_1106/train/test0513.ply')
                        #     trimesh.Scene([pc_1, pc, mesh]).show()
                        break
        return [root, joint]

    def refine(self, grasp_config, vis=False):
        '''
        grasp_config: [root, joint]
        '''        
        root:torch.FloatTensor = grasp_config[0]
        joint:torch.FloatTensor = grasp_config[1]
        t = torch.clone(root[:, :3])
        q = torch.clone(root[:, 3:])
        t.requires_grad_()
        # root.requires_grad_()
        joint.requires_grad_()
        bs = root.size(0)
        
        if vis:
            vis_o3d = o3d.visualization.Visualizer()
            vis_o3d.create_window()
            mesh = o3d_read_mesh('dif/recon/mug_1106/train/test0513.ply', scale=1/self.scale)
            vis_o3d.add_geometry(mesh)
            hand_mesh = self.fk.get_posed_mesh(root[0].detach().cpu().numpy(),
                                               joint[0].detach().cpu().numpy())
            shadow_v = o3d.utility.Vector3dVector(hand_mesh.vertices)
            shadow_f = o3d.utility.Vector3iVector(hand_mesh.faces)            
            hand_mesh = o3d.geometry.TriangleMesh(shadow_v, shadow_f)
            hand_mesh.compute_vertex_normals()
            vis_o3d.add_geometry(hand_mesh)
            shadow_pc = o3d.geometry.PointCloud()
            shadow_pc.points = o3d.utility.Vector3dVector(shadow_v)
            vis_o3d.add_geometry(shadow_pc)
        
        optim = torch.optim.Adamax([{'params': t, 'lr': 1e-4},
                                    {'params': joint, 'lr': 1e-2}])
        
        # >>>>>> select grasp points
        # with torch.no_grad():
        #     root = torch.cat([t, q], dim=1)
        #     out = self.fk.get_posed_pc(root, joint)
        #     pcs = out.verts
        #     norm_pcs = pcs * self.scale
            
        #     codes = torch.clone(self.code).repeat(bs, 1)
        #     pcs_template = self.dif.get_template_coords(norm_pcs, codes) # [1, hand pts num, 3]
        #     pcs_sdfs = self.dif.get_template_field(pcs_template) # [1, hand pts num, 1]
            
        #     inner_idxs = torch.norm(pcs_template, dim=2) < 0.99
        #     contact_idxs = pcs_sdfs[:,:,0] < 5e-3
        #     non_contact_idxs = pcs_sdfs[:,:,0] > 5e-2
        #     contact_idxs = contact_idxs & inner_idxs
        #     non_contact_idxs = non_contact_idxs & inner_idxs
        # <<<<<<
            
        for i in range(self.refine_steps):
            root = torch.cat([t, q], dim=1)
            out = self.fk.get_posed_pc(root, joint)
            pcs = out.verts
            norm_pcs = pcs * self.scale
            
            codes = torch.clone(self.code).repeat(bs, 1)
            pcs_template = self.dif.get_template_coords(norm_pcs, codes) # [1, hand pts num, 3]
            pcs_sdfs = self.dif.get_template_field(pcs_template) # [1, hand pts num, 1]
            
            inner_idxs = torch.norm(pcs_template, dim=2) < 0.99
            contact_idxs = pcs_sdfs[:,:,0] < 5e-2
            non_contact_idxs = pcs_sdfs[:,:,0] > 5e-1
            contact_idxs = contact_idxs & inner_idxs
            non_contact_idxs = non_contact_idxs & inner_idxs
                                
            l_fc = 0.0
            l_cont = 0.0
            l_colli = 0.0
            for j in range(bs):
                contact_idx = contact_idxs[j]
                non_contact_idx = non_contact_idxs[j]
                # >>>>>> compute FC loss
                contact_pcs = torch.clone(norm_pcs[j, contact_idx]).unsqueeze(0)
                contact_pcs_sdfs = self.dif.inference_with_grad(contact_pcs, codes[j])

                contact_pts_normals = torch.autograd.grad(contact_pcs_sdfs, contact_pcs, 
                                                        grad_outputs=torch.ones_like(contact_pcs_sdfs), 
                                                        create_graph=True)[0]
                normals = contact_pts_normals / torch.norm(contact_pts_normals, dim=2, keepdim=True)
                f = normals * contact_pcs_sdfs * self.scale
                l_fc_a, l_fc_b = self.fc.fc_loss(pcs[j, contact_idx].unsqueeze(0), f)
                l_fc_b = torch.where(l_fc_b > 0.5, l_fc_b, torch.zeros_like(l_fc_b))
                l_fc += l_fc_a + l_fc_b
                # <<<<<<
                
                # >>>>>> compute contact loss
                l_contact = torch.where(contact_pcs_sdfs > 0, contact_pcs_sdfs, torch.zeros_like(contact_pcs_sdfs))
                l_contact = l_contact.mean()
                l_cont += l_contact
                # <<<<<<
                
                # >>>>>> compute collision loss
                non_contact_pcs = torch.clone(norm_pcs[j, :]).unsqueeze(0)
                non_contact_pcs_sdfs = self.dif.inference_with_grad(non_contact_pcs, codes[j])        
                l_collision = -torch.where(non_contact_pcs_sdfs < 0, non_contact_pcs_sdfs, torch.zeros_like(non_contact_pcs_sdfs))
                l_collision = l_collision.mean()
                l_colli += l_collision
                # <<<<<<
                
            # >>>>>> compute limit loss
            l_lmt = self.fk.compute_loss_limit(joint)
            
            l_fc *= 1e-4
            l_cont *= 1e0
            l_colli *= 1e1
            l_lmt *= 1e1
            l = l_fc + l_cont + l_colli + l_lmt
            optim.zero_grad()
            l.backward(retain_graph=True)
            optim.step()
                                                  
            if i % 20 == 0 or i == self.refine_steps - 1:
                print(f'[{i:5d} / {self.refine_steps:5d}]' +
                      f'fc: {l_fc.item():.3f},' +
                      f'cont: {l_cont.item():.3f},' +
                      f'colli: {l_colli.item():.3f},' +
                      f'lmt: {l_lmt.item():.3f}')
                
            if vis:
                m = self.fk.get_posed_mesh(root[0].detach().cpu().numpy(),
                                            joint[0].detach().cpu().numpy())
                shadow_v = o3d.utility.Vector3dVector(m.vertices)                    
                hand_mesh.vertices = shadow_v
                shadow_pc.points = o3d.utility.Vector3dVector(pcs[0].detach().cpu().numpy())
                vis_o3d.update_geometry(hand_mesh)
                vis_o3d.update_geometry(shadow_pc)
                vis_o3d.poll_events()
                vis_o3d.update_renderer()                


                    
                
    
    def run(self, vis=False):
        grasp_list = self.get_init_grasps(num=1, start_idx=1)
        grasp_config = self.preprocess_grasps(grasp_list, backoff=False)
        grasp_config = self.refine(grasp_config, vis)

if __name__ == '__main__':
    cat = 'mug'
    device = 'cuda:0'
    init_path = f'{init_dir}/{cat}'
    runner = GraspSythesis(cat, init_path, device)
    runner.run(vis=True)